#!/usr/bin/env python
#-*- coding:utf-8 -*-


# import numpy as np
from common.node import Node, NodeSet
from common.edge import Edge, EdgeSet
from graph.graph import Graph
from displayer.displayer import PyplotDisplayer

class NodeGraph(Graph): 
    def __init__(self, nodes=None, edges=None):
        self.init(nodes, edges)

    def init(self, nodes=None, edges=None):
        self.node_set = nodes
        self.edge_set = edges

    @property 
    def node_set(self):
        return self._node_set
    
    @node_set.setter
    def node_set(self, node_set):
        if node_set is None:
            self._node_set = NodeSet()
        else:
            self._node_set = node_set
        
    @property 
    def edge_set(self):
        return self._edge_set
    
    @edge_set.setter
    def edge_set(self, edge_set):
        if edge_set is None:
            self._edge_set = EdgeSet()
        else:
            self._edge_set = edge_set

    @property
    def node_size(self):
        return len(self.node_set)

    @property
    def heuristics(self):
        return self._heuristics

    @heuristics.setter
    def heuristics(self, heuristics):
        if heuristics is None:
            self._heuristics = {}
        else:
            self._heuristics = heuristics

    def add_node(self, node):
        self.node_set.add_node(node)

    def add_edge(self, edge, update_succ=True, update_prec=False):
        self.edge_set.add_edge(edge)
        if update_succ:
            self.node_set.add_succ(edge.start, edge.end)
        if update_prec:
            self.node_set.add_prec(edge.end, edge.start)

    def node(self, nid):
        return self.node_set(nid)

    def nodes(self):
        return self.node_set.nodes()

    def position(self, nid):
        return self.node_set.position(nid)

    def weight(self, from_u, to_v):
        return self.edge_set.weight(from_u, to_v)

    def cost(self, from_u, to_v): # for interface
        return self.weight(from_u, to_v)

    def distance(self, from_u, to_v, *args, **kws):
        return self.weight(from_u, to_v)

    def blocked(self, from_u, to_v=None):
        blocks = []
        if not self.node(from_u):
            blocks.append(from_u)
        if to_v is None:
            return blocks
        if not self.node(to_v):
            blocks.append(to_v)
        if not self.weight(from_u, to_v):
            blocks.append(from_u)
            blocks.append(to_v)
        blocks = set(blocks)
        blocks = list(blocks)
        return blocks

    def successor(self, from_u):
        return self.node_set.succ(from_u)

    def heuristic(self, from_u, to_v):
        _key = (from_u, to_v)
        if _key in self.heuristics:
            return self.heuristics[_key]
        else:
            return None 

    def load(self, vertices, arcs, heuristics=None):
        self.node_set.reset()
        self.edge_set.reset()
        for nid, pos in vertices.items():
            self.add_node(Node(nid, pos))
        for arc, weight in arcs.items():
            self.add_edge(Edge(arc[0], arc[1], weight))
        self.heuristics = heuristics

    @property 
    def displayer(self):
        if not hasattr(self, "_displayer"):
            self._displayer = PyplotDisplayer()
        return self._displayer

    @displayer.setter
    def displayer(self, displayer):
        self._displayer = displayer

    def draw_nodes(self, color='blue', n_size=1000, annotate=False, trans = False): #magenta
        for node in self.node_set.nodes():
            x, y = node.pos
            x_d = x - 0.009 
            y_d = y - 0.015
            if trans: 
                x, y = y, x
                # x_d, y_d = y_d, x_d
            self.displayer.scatter(x, y, s=n_size, c=color, marker='o', alpha=0.8)
            if annotate:
                self.displayer.annotate(node.nid, (x_d, y_d), 
                    fontsize=20, color='white', multialignment='center', alpha=0.8)

    def draw_edges(self, color='blue', weight=False, trans=False): #'magenta'
        for edge in self.edge_set.edges():
            start_node = self.node_set(edge.start)
            end_node = self.node_set(edge.end)
            if not (start_node and end_node):
                return False
            x = [start_node.pos[0], end_node.pos[0]]
            y = [start_node.pos[1], end_node.pos[1]]
            x_c = (x[0] + x[1]) / 2 - 0.009
            y_c = (y[0] + y[1]) / 2 - 0.015
            if trans: 
                x, y = y, x
                x_c, y_c = y_c, x_c
            self.displayer.plot(x, y, color=color, linewidth=2, alpha=0.8)
            if weight:
                self.displayer.annotate(edge.weight, (x_c, y_c), 
                            fontsize=15, color='black', backgroundcolor='white', alpha=0.8)
            
    def draw(self, displayer=None, b_color='blue', n_size=1000, annotate=False, weight=False, trans=False):
        if displayer is not None:
            self.displayer = displayer
        self.draw_nodes(color=b_color, n_size=n_size, annotate=annotate, trans=trans)
        self.draw_edges(color=b_color, weight=weight, trans=trans)

    def savefig(self, fname=None, displayer=None):
        if displayer is not None:
            self.displayer = displayer
        self.displayer.savefig(fname=fname)


def test():
    vertices = {"S":(0.05, 0.5), #0 #含position信息
                "A":(0.3, 0.8), #1
                "B":(0.3, 0.5), #2
                "C":(0.3, 0.2), #3
                "D":(0.6, 0.95), #4
                "E":(0.6, 0.65), #5
                "F":(0.6, 0.4),  #6
                "K":(0.8, 0.2), #7
                "G":(0.99, 0.5),} #8
    arcs = { ('S', 'A'): 9, ('S', 'B'): 6, ('S', 'C'): 8, ('A', 'S'): 9, ('B', 'S'): 6, ('C', 'S'): 8, #弧是有向边，边可以去重
                ('A', 'D'): 7, ('A', 'E'): 9, ('D', 'A'): 7, ('E', 'A'): 9,
                ('B', 'E'): 8, ('B', 'F'): 8, ('E', 'B'): 8, ('F', 'B'): 8,
                ('C', 'K'): 20, ('K', 'C'): 20,
                ('D', 'G'): 16, ('G', 'D'): 16,
                ('E', 'G'): 13, ('G', 'E'): 13, 
                ('F', 'G'): 13, ('F', 'K'): 5, ('G', 'F'): 13, ('K', 'F'): 5, 
                ('K', 'G'): 6, ('G', 'K'): 6 }
    heuristics = {  ("S", "G"): 20, #0
                    ("A", "G"): 15, #1
                    ("B", "G"): 17, #2
                    ("C", "G"): 15, #3
                    ("D", "G"): 11, #4
                    ("E", "G"): 12, #5
                    ("F", "G"): 10, #6
                    ("K", "G"): 5,  #7
                    ("G", "G"): 0, } #8
    node_graph = NodeGraph()
    node_graph.load(vertices, arcs, heuristics)


if __name__ == "__main__":
    test()